from itertools import permutations, product
import numpy as np


def random_crop(iterable, crop_shape, seed):
    crop_shape = np.asarray(crop_shape, 'int')
    rng = np.random.RandomState(seed)

    def _random_crop(x):
        x = np.asarray(x)
        ori_shape = np.asarray(x.shape, 'int')
        diff = ori_shape - crop_shape
        assert np.all(diff >= 0), (ori_shape, crop_shape)
        starts = (rng.randint(d + 1) for d in diff)
        slices = tuple(slice(b, b+s) for b, s in zip(starts, crop_shape))
        cropped = x[slices]
        assert np.all(np.asarray(cropped.shape, 'int') == crop_shape)
        return cropped

    return map(_random_crop, iterable)


def middle_crop(iterable, crop_shape):
    crop_shape = np.asarray(crop_shape, 'int')

    def _middle_crop(x):
        x = np.asarray(x)
        ori_shape = np.asarray(x.shape, 'int')
        diff = ori_shape - crop_shape
        assert np.all(diff >= 0), (ori_shape, crop_shape)
        starts = (d // 2 for d in diff)
        slices = tuple(slice(b, b + s) for b, s in zip(starts, crop_shape))
        cropped = x[slices]
        assert np.all(np.asarray(cropped.shape, 'int') == crop_shape)
        return cropped

    return map(_middle_crop, iterable)


def vertical_rotation_and_flip(iterable, ndim, seed):
    axes = tuple(permutations(range(ndim)))
    steps = tuple(product((-1, 1), repeat=ndim))
    rng = np.random.RandomState(seed)

    def _vertical_rotation_and_flip(x):
        x = np.asarray(x)
        x = np.transpose(
            x,
            axes[rng.randint(len(axes))] + tuple(range(ndim, x.ndim))
        )
        x = x[
            tuple(slice(None, None, s) for s in steps[rng.randint(len(steps))])
        ]
        return x

    return map(_vertical_rotation_and_flip, iterable)


def random_flip(iterable, ndim, seed):
    steps = tuple(product((-1, 1), repeat=ndim))
    rng = np.random.RandomState(seed)

    def _random_flip(x):
        x = np.asarray(x)
        x = x[
            tuple(slice(None, None, s) for s in steps[rng.randint(len(steps))])
        ]
        return x

    return map(_random_flip, iterable)
